package cn.cxq.learning.tree;

import java.io.*;
import java.util.*;

/**
 * 哈夫曼编码
 */
public class HuffmanCode {

    private String huffmanCode; //哈夫曼编码后的字符串
    private HuffmanCodeTree huffmanCodeTree; //哈夫曼编码用到的哈夫曼树，节点的信息包括字符对应的ascii码以及权重（通过字符出现频率来描述）
    private Map<Byte, String> huffmanCodeTable; //哈夫曼字符ascii码和其编码的对应表
    private int lastLength; // 最后一个字节的位数
    private int[] deZero; // 将所有的负0记录下来

    public String getHuffmanCode() {
        return huffmanCode;
    }

    public void setHuffmanCode(String huffmanCode) {
        this.huffmanCode = huffmanCode;
    }

    public HuffmanCodeTree getHuffmanCodeTree() {
        return huffmanCodeTree;
    }

    public void setHuffmanCodeTree(HuffmanCodeTree huffmanCodeTree) {
        this.huffmanCodeTree = huffmanCodeTree;
    }

    public Map<Byte, String> getHuffmanCodeTable() {
        return huffmanCodeTable;
    }

    public void setHuffmanCodeTable(Map<Byte, String> huffmanCodeTable) {
        this.huffmanCodeTable = huffmanCodeTable;
    }

    /**
     * @param msg 待编码的字符串
     */
    public HuffmanCode(String msg) {

        this.huffmanCodeTree = new HuffmanCodeTree(msg);
        this.huffmanCodeTable = this.huffmanCodeTree.initHuffmanCodeTable();


        this.huffmanCode = initHuffmanCode(msg);
    }

    public HuffmanCode() {
    }

    /**
     * 生成消息字符串的哈夫曼编码
     *
     * @param msg 原文
     * @return 哈夫曼编码
     */
    private String initHuffmanCode(String msg) {

        StringBuffer buffer = new StringBuffer();

        for (int i = 0; i < msg.length(); i++) {
            byte ascii = (byte) msg.charAt(i);
            for (Byte c : huffmanCodeTable.keySet()) {
                if (ascii == c) {
                    buffer.append(huffmanCodeTable.get(c));
                    break;
                }
            }
        }

        return buffer.toString();
    }

    private String initHuffmanCode(byte[] bytes) {

        StringBuffer buffer = new StringBuffer();

        for (int i = 0; i < bytes.length; i++) {
            for (Byte c : huffmanCodeTable.keySet()) {
                if (bytes[i] == c) {
                    buffer.append(huffmanCodeTable.get(c));
                    break;
                }
            }
        }

        return buffer.toString();
    }

    public byte[] getHuffmanBytes() {
        int length;
        if (huffmanCode.length() % 8 == 0) length = huffmanCode.length() / 8;
        else length = huffmanCode.length() / 8 + 1;

        byte[] huffmanBytes = new byte[length];

        int index = 0;

        List<Integer> deZeroList = new ArrayList<>();

        for (int i = 0; i < huffmanCode.length(); i += 8) {

            String value;

            if (i + 8 < huffmanCode.length()) value = huffmanCode.substring(i, i + 8);
            else {
                value = huffmanCode.substring(i);
                this.lastLength = value.length();
            }

            //如果是负数则需要求补码
            if (value.length() == 8 && value.charAt(0) == '1') {
                int j;
                for (j = value.length() - 1; j >= 0; j--) {
                    if (value.charAt(j) == '1') break;
                }

                j--;

                while (j > 0) {
                    if (value.charAt(j) == '0') value = value.substring(0, j) + '1' + value.substring(j + 1);
                    else value = value.substring(0, j) + '0' + value.substring(j + 1);
                    j--;
                }
                if (value.equals("10000000")) deZeroList.add(index);
                value = '-' + value.substring(1);
            }

            huffmanBytes[index++] = Byte.parseByte(value, 2);
        }

        this.deZero = new int[deZeroList.size()];

        for (int i = 0; i < deZeroList.size(); i++) {
            this.deZero[i] = deZeroList.get(i);
        }

        return huffmanBytes;
    }

    /**
     * 哈夫曼编码解码， 通过哈夫曼编码和哈夫曼编码表实现
     *
     * @return 原文
     */
    public String decode() {
        int begin = 0;
        StringBuffer buffer = new StringBuffer();
        for (int i = 1; i <= this.huffmanCode.length(); i++) {
            String code = this.huffmanCode.substring(begin, i);
            for (Byte c : huffmanCodeTable.keySet()) {
                if (huffmanCodeTable.get(c).equals(code)) {
                    begin = i;
                    buffer.append((char) (byte) c);
                    break;
                }
            }
        }

        return buffer.toString();
    }

    public byte[] decode(byte[] huffmanCodeBytes, Map<Byte, String> huffmanCodeTable, int lastLength, int[] deZero) {

        List<Byte> byteList = new ArrayList<>();

        StringBuffer buffer = new StringBuffer();
        for (int i = 0; i < huffmanCodeBytes.length; i++) {
            String temp;
            if (i != huffmanCodeBytes.length - 1) temp = byteToString(huffmanCodeBytes[i], 8);
            else temp = byteToString(huffmanCodeBytes[i], lastLength);

            for (int j = 0; j < deZero.length; j++) {
                if (i == deZero[j]) {
                    temp = '1' + temp.substring(1);
                    break;
                }
            }

            buffer.append(temp);
        }

        String code = buffer.toString();

//        int num = 0;
//
//        for (int i = 0; i < code.length(); i++) {
//            if (this.huffmanCode.charAt(i) != code.charAt(i)) {
//                System.out.println(i+"位:"+this.huffmanCode.charAt(i)+ ":" + code.charAt(i));
//                num ++;
//            }
//        }
//
//        System.out.println(num);

        int begin = 0;

        for (int i = 1; i <= code.length(); i++) {
            String temp = code.substring(begin, i);
            for (Byte c : huffmanCodeTable.keySet()) {
                if (huffmanCodeTable.get(c).equals(temp)) {
                    begin = i;
                    byteList.add(c);
                    break;
                }
            }
        }

        byte[] bytes = new byte[byteList.size()];

        int index = 0;

        for (Byte aByte : byteList) {
            bytes[index++] = aByte;
        }

        return bytes;
    }

    /**
     * 将字节转换成二进制字符串
     *
     * @param b
     * @param length
     * @return
     */
    private String byteToString(byte b, int length) {

        int temp = b;
        temp |= 256;
        String value = Integer.toBinaryString(temp);
        value = value.substring(value.length() - length);

        return value;
    }

    public void zipFile(String src, String dst) {

        InputStream is = null;
        OutputStream os = null;
        ObjectOutputStream oos = null; //创建对象流

        try {
            is = new FileInputStream(src);

            //创建一个和源文件大小一样的byte数组
            byte[] bytes = new byte[is.available()];
            is.read(bytes);

            //通过字节数组获取哈夫曼树
            this.huffmanCodeTree = new HuffmanCodeTree(bytes);

            //通过哈夫曼树获取哈夫曼编码表
            this.huffmanCodeTable = this.huffmanCodeTree.initHuffmanCodeTable();

            //获取哈夫曼编码
            this.huffmanCode = initHuffmanCode(bytes);

            //将哈夫曼编码以每8位（二进制转十进制）为一个单位存储到一个字节数组中
            byte[] huffmanBytes = getHuffmanBytes();

            os = new FileOutputStream(dst);
            oos = new ObjectOutputStream(os);

            //将哈夫曼编码写入目标文件中
            oos.writeObject(huffmanBytes);

            //将哈夫曼编码表写入目标文件中
            oos.writeObject(this.huffmanCodeTable);

            //将最后一字节的位数传入
            oos.writeObject(this.lastLength);

            //将所有负零所在的坐标传入
            oos.writeObject(this.deZero);

        } catch (FileNotFoundException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                if (is != null) is.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
            try {
                if (oos != null) oos.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
            try {
                if (os != null) os.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    public void unZipFile(String src, String dst) {

        InputStream is = null;
        OutputStream os = null;
        ObjectInputStream ois = null; //创建对象输入流

        try {
            is = new FileInputStream(src);
            ois = new ObjectInputStream(is);
            byte[] huffmanBytes = (byte[]) ois.readObject(); //获取哈夫曼编码

            Map<Byte, String> huffmanCodeTable = (Map<Byte, String>) ois.readObject(); //获取哈夫曼编码表


            int lastLength = (int) ois.readObject(); //获取最后一字节的位数

            int[] deZero = (int[]) ois.readObject(); //获取所有的负零坐标

            byte[] bytes = decode(huffmanBytes, huffmanCodeTable, lastLength, deZero);

            os = new FileOutputStream(dst);

            os.write(bytes);

        } catch (FileNotFoundException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (ClassNotFoundException e) {
            e.printStackTrace();
        } finally {
            try {
                if (is != null) is.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
            try {
                if (ois != null) ois.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
            try {
                if (os != null) os.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    public int getLastLength() {
        return lastLength;
    }

    public void setLastLength(int lastLength) {
        this.lastLength = lastLength;
    }
}

class HuffmanCodeTree {

    private HuffmanCodeNode root;

    public HuffmanCodeTree(byte[] bytes) {

        HashMap<Byte, Integer> charHashMap = new HashMap<>();

        for (int i = 0; i < bytes.length; i++) {
            if (charHashMap.containsKey(bytes[i])) charHashMap.put(bytes[i], charHashMap.get(bytes[i]) + 1);
            else charHashMap.put(bytes[i], 1);
        }

        List<HuffmanCodeNode> huffmanNodeList = new ArrayList<>();

        for (Byte c : charHashMap.keySet()) {

            HuffmanCodeNode huffmanNode = new HuffmanCodeNode(charHashMap.get(c), c);

            huffmanNodeList.add(huffmanNode);
        }


        while (huffmanNodeList.size() > 1) {

            Collections.sort(huffmanNodeList);

            HuffmanCodeNode leftNode = huffmanNodeList.get(0);
            HuffmanCodeNode rightNode = huffmanNodeList.get(1);

            HuffmanCodeNode newNode = new HuffmanCodeNode(leftNode.weight + rightNode.weight, null);

            newNode.left = leftNode;
            newNode.right = rightNode;

            huffmanNodeList.remove(leftNode);
            huffmanNodeList.remove(rightNode);
            huffmanNodeList.add(newNode);

        }

        this.root = huffmanNodeList.get(0);
    }

    public HuffmanCodeTree(String msg) {

        HashMap<Byte, Integer> charHashMap = new HashMap<>();

        for (int i = 0; i < msg.length(); i++) {

            char c = msg.charAt(i);

            //注意这里的键值并没有做限制，所以要注意将类型强转成byte
            if (charHashMap.containsKey((byte) c)) charHashMap.put((byte) c, charHashMap.get((byte) c) + 1);

            else charHashMap.put((byte) c, 1);
        }

        List<HuffmanCodeNode> huffmanNodeList = new ArrayList<>();

        for (Byte c : charHashMap.keySet()) {

            HuffmanCodeNode huffmanNode = new HuffmanCodeNode(charHashMap.get(c), c);

            huffmanNodeList.add(huffmanNode);
        }


        while (huffmanNodeList.size() > 1) {

            Collections.sort(huffmanNodeList);

            HuffmanCodeNode leftNode = huffmanNodeList.get(0);
            HuffmanCodeNode rightNode = huffmanNodeList.get(1);

            HuffmanCodeNode newNode = new HuffmanCodeNode(leftNode.weight + rightNode.weight, null);

            newNode.left = leftNode;
            newNode.right = rightNode;

            huffmanNodeList.remove(leftNode);
            huffmanNodeList.remove(rightNode);
            huffmanNodeList.add(newNode);

        }

        this.root = huffmanNodeList.get(0);
    }

    public void preorderTransversal() {
        if (this.root == null) System.out.println("哈夫曼树为空~~~");
        else this.root.preorderTransversal();
    }

    /**
     * 初始化哈夫曼编码表
     *
     * @return
     */
    public Map<Byte, String> initHuffmanCodeTable() {
        Map<Byte, String> huffmanCodeTable = null;
        if (this.root == null) System.out.println("哈夫曼树为空~~~");
        else {
            huffmanCodeTable = new HashMap<>();
            if (this.root.left != null) this.root.initHuffmanCode(huffmanCodeTable, 0, new StringBuffer());
            if (this.root.right != null) this.root.initHuffmanCode(huffmanCodeTable, 1, new StringBuffer());
        }
        return huffmanCodeTable;
    }
}

class HuffmanCodeNode implements Comparable<HuffmanCodeNode> {

    int weight;
    Byte c; //字符串对应的ascii码
    HuffmanCodeNode left;
    HuffmanCodeNode right;

    public HuffmanCodeNode(int value, Byte c) {
        this.weight = value;
        this.c = c;
    }

    @Override
    public String toString() {
        return "HuffmanCodeNode{" +
                "weight=" + weight +
                ", c=" + c +
                '}';
    }

    @Override
    public int compareTo(HuffmanCodeNode huffmanNode) {
        return this.weight - huffmanNode.weight;
    }

    public void preorderTransversal() {

        System.out.println(this);

        if (this.left != null) this.left.preorderTransversal();
        if (this.right != null) this.right.preorderTransversal();
    }

    /**
     * 递归生成哈夫曼编码存到哈夫曼编码表中
     *
     * @param huffmanCodeTable
     * @param i
     * @param stringBuffer
     */
    public void initHuffmanCode(Map<Byte, String> huffmanCodeTable, int i, StringBuffer stringBuffer) {

        StringBuffer buffer = new StringBuffer(stringBuffer);

        if (i == 0) {
            buffer.append(0);
            HuffmanCodeNode left = this.left;
            if (left.left != null) {
                this.left.initHuffmanCode(huffmanCodeTable, 0, buffer);
            }
            if (left.right != null) {
                this.left.initHuffmanCode(huffmanCodeTable, 1, buffer);
            }
            if (left.left == null && left.right == null) huffmanCodeTable.put(left.c, buffer.toString());
        } else {
            buffer.append(1);
            HuffmanCodeNode right = this.right;
            if (right.left != null) {
                this.right.initHuffmanCode(huffmanCodeTable, 0, buffer);
            }
            if (right.right != null) {
                this.right.initHuffmanCode(huffmanCodeTable, 1, buffer);
            }
            if (right.left == null && right.right == null) huffmanCodeTable.put(right.c, buffer.toString());
        }

    }
}
